### graph 
import numpy as np
import generators as G

#TODO:
#  let constructors specify n=__, edges=__, matrix=__, symmetric=True/False storage=AdjacencyList/AdjacencyMatrix
#  and have just a single Graph class

class VertexListImplementation(object):
	def __init__(self,n=0):
		self._v=[None for _ in range(n)]
	
	def get_vertex(self,v):
		return self._v[v]
	
	def set_vertex(self,v,x):
		self._v[v]=x
	
	# del_vertex -- TODO, messy and not crucial
	
	def add_vertex(self,x=None):
		rv=len(self._v)
		self._v.append(x)
		return rv

	def num_vertices(self):
		return len(self._v)

class VertexNullImplementation(object):
	def __init__(self,n=0):
		self._n=n
		
	def get_vertex(self,v):
		return None
	
	def set_vertex(self,v,x):
		pass
	
	def add_vertex(self,x=None):
		rv=self._n
		self._n+=1
		return rv
	
	def num_vertices(self):
		return self._n

class AdjacencyListImplementation(object):
	def __init__(self,n=0):
		self._e=[{} for _ in xrange(n)]
		self._m=0

	def neighbors(self,u):
		for v in self._e[u]:
			yield v
	
	def get_edge(self,u,v):
		if v in self._e[u]:
			return self._e[u][v]
		else:
			return None
	
	def set_edge(self,u,v,x):
		if v not in self._e[u]:
			self._m+=1
		self._e[u][v]=x

	def del_edge(self,u,v):
		if v in self._e[u]:
			self._m-=1
		del self._e[u][v]
	
	def add_vertex(self):
		self._e.append({})
	
	def num_edges(self):
		return self._m

	def __repr__(self):
		return '%s(%s)' % (self.__class__.__name__,repr(self._e))

	def __str__(self):
		return str(self._e)
		

class AdjacencyMatrixImplementation(object):
	def __init__(self,n=0):
		self._e = np.empty((n,n),object)
		self._m = 0

	def neighbors(self,u):
		for v in xrange(self._e.shape[0]):
			if self._e[u,v]!=None:
				yield v

	def get_edge(self,u,v):
		return self._e[u,v]

	def set_edge(self,u,v,x):
		if self._e[u,v] is None:
			self._m+=1
		self._e[u,v]=x

	def del_edge(self,u,v):
		if self._e[u,v] is not None:
			self._m-=1
		self._e[u,v]=None

	def add_vertex(self): # slow!!!
		old_e=self._e
		old_n=old_e.shape[0]
		self._e=np.empty((old_n+1,old_n+1),object)
		self._e[:old_n,:old_n]=old_e

	def num_edges(self):
		return self._m

	def __repr__(self):
		return '%s(%s)' % (self.__class__.__name__,repr(self._e))

	def __str__(self):
		return str(self._e)


class _EdgesProxy(object):
	def __init__(self,graph):
		self._graph=graph
	
	def __getitem__(self,(u,v)):
		return self._graph.get_edge(u,v)
	
	def __setitem__(self,(u,v),x):
		self._graph.set_edge(u,v,x)
	
	def __delitem__(self,(u,v)):
		return self._graph.del_edge(u,v)

	def __len__(self):
		return self._graph.num_edges()

	def __iter__(self):
		for u in self._graph.vertices:
			for v in self._graph.neighbors(u):
				yield (u,v)


class _VerticesProxy(object):
	def __init__(self,graph):
		self._graph=graph
	
	def __getitem__(self,u):
		return self._graph.get_vertex(u)
	
	def __setitem__(self,u,x):
		self._graph.set_vertex(u,x)
	
	def __delitem__(self,u):
		self._graph.del_vertex(u)
	
	def __len__(self):
		return self._graph.num_vertices()

	def __iter__(self):
		return iter(xrange(self._graph.num_vertices()))


class Graph(object):
	def __init__(self,graph=None, n=0, edges=None, matrix=None, symmetric=False, 
	             implementation=AdjacencyListImplementation, vertex_implementation=VertexNullImplementation,
	             is_edge=(lambda x: x), default_edge=1):
		self.edges=_EdgesProxy(self)
		self.vertices=_VerticesProxy(self)
		if edges:
			if matrix:
				raise ValueError("At most one of edges and matrix may be specified")
			n=max(n,1+max(max(u,v) for (u,v) in edges))
		elif matrix:
			try:
				n=max(n,max(matrix.shape))
			except:
				n=max(n,max(len(matrix),max(len(r) for r in matrix)))
		self._impl=implementation(n)
		self._vimpl=vertex_implementation(n)
		self._symmetric=symmetric
		if symmetric:
			try:
				symmetric.__call__
			except:
				self._symmetric=lambda x: x
			self.get_edge=self.get_edge_symmetric
			self.set_edge=self.set_edge_symmetric
			self.del_edge=self.del_edge_symmetric
		else:
			self.get_edge=self.get_edge_asymmetric
			self.set_edge=self.set_edge_asymmetric
			self.del_edge=self.del_edge_asymmetric
		
		self.is_edge=is_edge
		self.default_edge=default_edge
		
		if edges:
			try:
				for (u,v) in edges:
					self.edges[u,v]=edges[u,v]
			except:
				for (u,v) in edges:
					self.edges[u,v]=default_edge
		elif matrix:
			for (u,r) in G.xzip(G.iota(),matrix):
				for (v,x) in G.xzip(G.iota(),r):
					if is_edge(x):
						self.edges[u,v]=x
		
	# convert to a numpy array
	def __array__(self,dtype=object):
		if dtype==object:
			rv=np.empty((self.num_vertices(),self.num_vertices()),object)
		else:
			rv=np.zeros((self.num_vertices(),self.num_vertices()))
		for (u,v) in self.edges:
			rv[u,v]=self.edges[u,v]
		return rv
	
	def clear(self):
		for u in self.vertices:
			for v in self.vertices:
				del self.edges[u,v]
				
	def cycle(self):
		for i in self.vertices:
			self.edges[i,(i+1)%self.num_vertices()]=self.default_edge
	
	def complete(self,*parts):
		if len(parts)==0:
			parts=(self.num_vertices(),)
		elif sum(parts)!=self.num_vertices():
			raise ValueError
		if len(parts)==1:
			for u in self.vertices:
				for v in self.vertices:
					self.edges[u,v]=self.default_edge
		else:
			offsets=[]
			o=0
			for p in parts:
				offsets.append(o)
				o+=p
			for ((pu,ou),(pv,ov)) in G.choose(zip(parts,offsets),2):
				for u in xrange(ou,ou+pu):
					for v in xrange(ov,ov+pv):
						self[u,v]=self.default_edge
						self[v,u]=self.default_edge
						
	
	# single source shortest path (Dijkstra)
	# does not work if value(u,v) is negative!!
	# returns a pair of lists sssp(u)=(dist to v, parent of v)
	def sssp(self,u):
		import heapq		
		dist=[None]*self.num_vertices()
		parent=[None]*self.num_vertices()
		dist[u]=0
		q=[(0,u)]
		done=set()
		while q!=[]:
			(dux,x)=heapq.heappop(q)
			if x in done: continue
			done.add(x)
			for y in self.neighbors(x):
				dxy=self.edges[x,y]
				if dist[y]==None or dist[y]>dux+dxy:
					duy=dux+dxy
					dist[y]=duy
					parent[y]=x
					heapq.heappush(q,(duy,y))
		return (dist,parent)
	
	# all pairs shortest path (Floyd-Warshall)
	def apsp(self):
		dist=np.array(self)
		parent=np.empty((self.num_vertices(),self.num_vertices()),object)
		for u in self.vertices:
			dist[u,u]=0
			for v in self.vertices:
				parent[v,u]=v
			parent[u,u]=None
		for k in self.vertices:
			for i in self.vertices:
				if dist[i,k]==None: continue
				for j in self.vertices:
					if dist[k,j]==None: continue
					if dist[i,j]==None or dist[i,j] > dist[i,k]+dist[k,j]:
						dist[i,j]=dist[i,k]+dist[k,j]
						parent[i,j]=parent[k,j]
		return (dist,parent)
	
	# breadth first traversal starting at u
	# for each edge, update x=f(u,v,x)
	def bfs_reduce(self,f,u,x):
		done=set([u])
		q=[u]
		while q!=[]:
			u=q.pop(0)
			for v in self.neighbors(u):
				if v in done: continue
				done.add(v)
				x=f(u,v,x)
				q.append(v)
		return x
	
	# depth first traversal starting at u
	# for each edge, update x=f(u,v,x)
	def dfs_reduce(self,f,u,x):
		done=set([u])
		def go(u,x):
			for v in self.neighbors(u):
				if v in done: continue
				done.add(v)
				x=f(u,v,x)
				x=go(v,x)
			return x
		return go(u,x)
	
	### Graph implementation glue
	def num_edges(self):
		return self._impl.num_edges()
		
	def get_edge_asymmetric(self,u,v):
		return self._impl.get_edge(u,v)
		
	def get_edge_symmetric(self,u,v):
		return self._impl.get_edge(u,v)

	def set_edge_asymmetric(self,u,v,x):
		if self.is_edge(x):
			self._impl.set_edge(u,v,x)
		else:
			self._impl.del_edge(u,v)
			
	def set_edge_symmetric(self,u,v,x):
		if self.is_edge(x):
			self._impl.set_edge(u,v,x)
			self._impl.set_edge(v,u,self._symmetric(x))
		else:
			self._impl.del_edge(u,v)
			self._impl.del_edge(v,u)
	
	def del_edge_asymmetric(self,u,v):
		self._impl.del_edge(u,v)
	
	def del_edge_symmetric(self,u,v):
		self._impl.del_edge(u,v)
		self._impl.del_edge(v,u)
	
	def neighbors(self,u):
		return self._impl.neighbors(u)
	
	
	def num_vertices(self):
		return self._vimpl.num_vertices()
	
	def get_vertex(self,u):
		return self._vimpl.get_vertex(u)
	
	def set_vertex(self,u,x):
		self._vimpl.set_vertex(u,x)
	
	def add_vertex(self,x=None):
		self._impl.add_vertex()
		return self._vimpl.add_vertex(x)
		